import numpy as np
from sklearn.decomposition import SparseCoder
import os
import timeit

COMPONENT_COUNT = 100
PATCH_COUNT = 10000
EPOCH_COUNT = 16
SPARSE_CODE_COUNT = 800
SPARSITY_PARAMETER = 4.0
NON_NEGATIVE = True
IN_DIR = 'scbasis/'
OUT_FOLDER = 'sccodes/'

start = timeit.default_timer()

###
#Sets of stimuli to use to name the output sparse coding responses in a .npy:

#Imagenet
codesFile = OUT_FOLDER + 'imnetcodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

#Mod index: texture naturalistic stimuli
#codesFile = OUT_FOLDER + 'natcodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

#Mod index: texture noise stimuli
#codesFile = OUT_FOLDER + 'noisecodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

#Figure-ground
#codesFile = OUT_FOLDER + 'fgcodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

#Texture
#codesFile = OUT_FOLDER + 'texcodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

#Line angle stimuli
#codesFile = OUT_FOLDER + 'linecodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

###
#Load the PCA transformed V1 complex mean subtracted data for the dataset here

#Default is the original ImageNet input patches
pcaTransformed = np.load('imnet/pcaTransformed.npy')[:, :COMPONENT_COUNT]

#Mod index: texture naturalistic stimuli
#pcaTransformed = np.load('modindex_nat/pcaTransformed.npy')[:, :COMPONENT_COUNT]

#Mod index: texture noise stimuli
#pcaTransformed = np.load('modindex_noise/pcaTransformed.npy')[:, :COMPONENT_COUNT]

#Figure-ground stimuli
#pcaTransformed = np.load('fg/pcaTransformed.npy')[:, :COMPONENT_COUNT]

#Texture stimuli
#pcaTransformed = np.load('tex/pcaTransformed.npy')[:, :COMPONENT_COUNT]

#Line angle stimuli
#pcaTransformed = np.load('line/pcaTransformed.npy')[:, :COMPONENT_COUNT]

inputFile = IN_DIR + 'basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
print(inputFile)
print(codesFile)

basis = np.load(inputFile)

sc = SparseCoder(basis, transform_algorithm = 'lasso_cd', transform_alpha = SPARSITY_PARAMETER, positive_code = NON_NEGATIVE)

codes = sc.transform(pcaTransformed)

if not os.path.exists(OUT_FOLDER):
    os.makedirs(OUT_FOLDER)

np.save(codesFile, codes)

stop = timeit.default_timer()

print('Time: ', stop - start)
